'''
Created on 16 Nov 2010

@author: tiberious

Tag matroska files
'''


from ebml import EbmlReader, EbmlWriter, BinaryData
import sys

SINT, UINT, FLOAT, STRING, UTF8, DATE, MASTER, BINARY = range(8)

# Interesting Matroska tags, split into sections for easier tag reading.
# Tags not defined here are skipped while parsing.

MatroskaTagsSeek = {
    # Segment
    0x08538067: ('Segment', MASTER),
    # Seek information
    0x014D9B74: ('SeekHead', MASTER),
    0x0DBB: ('Seek', MASTER),
    0x13AB: ('SeekID', BINARY),
    0x13AC: ('SeekPosition', UINT),
    0x6C: ('EbmlVoid', BINARY)
}

MatroskaSeekIndex = {
    'Segment':      (0x08538067, MASTER),
    'SeekHead':     (0x014D9B74, MASTER),
    'Seek':         (0x0DBB, MASTER),
    'SeekID':       (0x13AB, BINARY),
    'SeekPosition': (0x13AC, UINT),
    'EbmlVoid':     (0x6C, BINARY)
}

MatroskaTagsSeekQuit = {
    0x0549A966: ('Info', MASTER),
}

MatroskaTagsInfo = {
    # Segment
    0x08538067: ('Segment', MASTER),                    
    # Segment Information
    0x0549A966: ('Info', MASTER),
    0x3384: ('SegmentFilename', UTF8),
    0x0AD7B1: ('TimecodeScale', UINT),
    0x0489: ('Duration', FLOAT),
    0x0461: ('DateUTC', DATE),
    0x3BA9: ('Title', UTF8),
    0x0D80: ('MuxingApp', UTF8),
    0x1741: ('WritingApp', UTF8),
    # Track
    0x0654AE6B: ('Tracks', MASTER),
    0x2E: ('TrackEntry', MASTER),
    0x57: ('TrackNumber', UINT),
    0x33C5: ('TrackUID', UINT),
    0x03: ('TrackType', UINT),
    0x29: ('FlagEnabled', UINT),
    0x08: ('FlagDefault', UINT),
    0x03E383: ('DefaultDuration', UINT),
    0x03314F: ('TrackTimecodeScale', FLOAT),
    0x137F: ('TrackOffset', SINT),
    0x136E: ('Name', UTF8),
    0x02B59C: ('Language', STRING),
    0x06: ('CodecID', STRING),
    0x058688: ('CodecName', UTF8),
    0x1A9697: ('CodecSettings', UTF8),
    0x1B4040: ('CodecInfoURL', STRING),
    0x06B240: ('CodecDownloadURL', STRING),
    0x2A: ('CodecDecodeAll', UINT),
    0x2FAB: ('TrackOverlay', UINT),
    # Video
    0x60: ('Video', MASTER),
    # Audio
    0x61: ('Audio', MASTER),
    0x35: ('SamplingFrequency', UINT),
    0x38B5: ('OutputSamplingFrequency', UINT),
    0x1F: ('Channels', UINT),
    0x3D7B: ('ChannelPositions', BINARY),
    0x2264: ('BitDepth', UINT),
    # Content Encoding
    0x2D80: ('ContentEncodings', MASTER),
    0x2240: ('ContentEncoding', MASTER),
    0x1031: ('ContentEncodingOrder', UINT),
    0x1032: ('ContentEncodingScope', UINT),
    0x1033: ('ContentEncodingType', UINT),
    0x1034: ('ContentCompression', MASTER),
    0x0254: ('ContentCompAlgo', UINT),
    0x0255: ('ContentCompSettings', BINARY),
    # Chapters
    0x0043A770: ('Chapters', MASTER),
    0x05B9: ('EditionEntry', MASTER),
    0x05BC: ('EditionUID', UINT),
    0x05BD: ('EditionFlagHidden', UINT),
    0x05DB: ('EditionFlagDefault', UINT),
    0x05DD: ('EditionManaged', UINT),
    0x36: ('ChapterAtom', MASTER),
    0x33C4: ('ChapterUID', UINT),
    0x11: ('ChapterTimeStart', UINT),
    0x12: ('ChapterTimeEnd', UINT),
    0x18: ('ChapterFlagHidden', UINT),
    0x0598: ('ChapterFlagEnabled', UINT),
    0x23C3: ('ChapterPhysicalEquiv', UINT),
    0x0F: ('ChapterTrack', MASTER),
    0x09: ('ChapterTrackNumber', UINT),
    0x00: ('ChapterDisplay', MASTER),
    0x05: ('ChapString', UTF8),
    0x037C: ('ChapLanguage', STRING),
    0x037E: ('ChapCountry', STRING),
    # Attachments
    0x0941A469: ('Attachments', MASTER),
    0x21A7: ('AttachedFile', MASTER),
    0x067E: ('FileDescription', UTF8),
    0x066E: ('FileName', UTF8),
    0x0660: ('FileMimeType', STRING),
    0x065C: ('FileData', BINARY),
    0x06AE: ('AttachmentUID', UINT)
}

MatroskaTagsInfoQuit = {
    0x0F43B675: ('Cluster', MASTER),
}

MatroskaTagsTagging = {
    # Segment
    0x08538067: ('Segment', MASTER),
    # Tagging
    0x0254C367: ('Tags', MASTER),
    0x3373: ('Tag', MASTER),
    0x23C0: ('Targets', MASTER),
    0x28CA: ('TargetTypevalue', UINT),
    0x23CA: ('TargetType', STRING),
    0x23C9: ('EditionUID', UINT),
    0x23C4: ('ChapterUID', UINT),
    0x23C5: ('TrackUID', UINT),
    0x23C6: ('AttachmentUID', UINT),
    0x27C8: ('SimpleTag', MASTER),
    0x05A3: ('TagName', UTF8),
    0x047A: ('TagLanguage', STRING),
    0x0484: ('TagDefault', UINT),
    0x0487: ('TagString', UTF8),
    0x0485: ('TagBinary', BINARY),
    # Seek information
    0x014D9B74: ('SeekHead', MASTER),
    0x0DBB: ('Seek', MASTER),
    0x13AB: ('SeekID', BINARY),
    0x13AC: ('SeekPosition', UINT),
    0x6C: ('EbmlVoid', BINARY)    
}

MatroskaTagsPosSeek = {
    # Segment
    0x08538067: ('Segment', MASTER),                         
    # Tagging
    0x0254C367: ('Tags', MASTER),
}

MatroskaTagsPosQuit = {
    0x3373: ('Tag', MASTER),
}

MatroskaTagsIndex = {
    'Segment':          (0x08538067, MASTER),
    # Tags
    'Tags':             (0x0254C367, MASTER),
    'Tag':              (0x3373, MASTER),
    'Targets':          (0x23C0, MASTER),
    'TargetTypevalue':  (0x28CA, UINT),
    'TargetType':       (0x23CA, STRING),
    'EditionUID':       (0x23C9, UINT),
    'ChapterUID':       (0x23C4, UINT),
    'TrackUID':         (0x23C5, UINT),
    'AttachmentUID':    (0x23C6, UINT),
    'SimpleTag':        (0x27C8, MASTER),
    'TagName':          (0x05A3, UTF8),
    'TagLanguage':      (0x047A, STRING),
    'TagDefault':       (0x0484, UINT),
    'TagString':        (0x0487, UTF8),
    'TagBinary':        (0x0485, BINARY),
    # Seek
    'SeekHead':         (0x014D9B74, MASTER),
    'Seek':             (0x0DBB, MASTER),
    'SeekID':           (0x13AB, BINARY),
    'SeekPosition':     (0x13AC, UINT),
    'EbmlVoid':         (0x6C, BINARY)    
}

MatroskaTagValidity = {
    # name: (isMulti, [Parents])
    'Tags':             (True, []),
    'Tag':              (True, ['Tags']),
    'Targets':          (False, ['Tag']),
    'TargetTypeValue':  (False, ['Targets']),
    'TargetType':       (False, ['Targets']),
    'EditionUID':       (True, ['Targets']),
    'ChapterUID':       (True, ['Targets']),
    'TrackUID':         (True, ['Targets']),
    'AttachmentUID':    (True, ['Targets']),
    'SimpleTag':        (True, ['Tag', 'SimpleTag']),
    'TagName':          (False, ['SimpleTag']),
    'TagLanguage':      (False, ['SimpleTag']),
    'TagDefault':       (False, ['SimpleTag']),
    'TagString':        (False, ['SimpleTag']),
    'TagBinary':        (False, ['SimpleTag']),
    }

class MatroskaTagReader:
    """
    Ebml Matroska Tag reader
    
    usage: MatroskaTagReader(location).getInfo()
    usage: MatroskaTagReader(location).getTags()
    """
    
    def __init__(self, location):
        self.ebml = EbmlReader(location, {})
        self.segmentPosition = -1
        self.info = None
        
    def __del__(self):
        del self.ebml
    
    def getInfo(self):
        if self.info == None:
            self.ebml.tags = MatroskaTagsInfo
            self.ebml.quitOn = MatroskaTagsInfoQuit
            self.ebml.storePosition = False
            try:
                self.info = self.ebml.parse()['Segment'][0]
            except (KeyError):
                return False               

        try:
            return self.info['Info'][0]
        except (KeyError):
            return False
        
    def getAllInfo(self):
        if self.info == None:
            self.ebml.tags = MatroskaTagsInfo
            self.ebml.quitOn = MatroskaTagsInfoQuit
            self.ebml.storePosition = False
            try:
                self.info = self.ebml.parse()['Segment'][0]
            except (KeyError):
                return False
            
        try:
            return self.info
        except:
            return False
        
    def _filterInfo(self, name, info, r):
        for i in info:
            try:
                for n, v in i.items():
                    if n in name:
                        try:
                            a = r[n]
                        except:
                            a = r[n] = []
                        a.append(v[0])
                    else:
                        self._filterInfo(name, v, r)
            except:
                return
        
    def getUIDs(self, names = None):
        if self.info == None:
            self.getAllInfo()
            
        # check it actually worked
        if self.info == None:
            return False
        
        if names == None:
            names = ['TrackUID', 'ChapterUID', 'EditionUID', 'AttachmentUID']
        
        # find all the relevant UIDs
        r = {}
        self._filterInfo(names, [self.info], r)
        
        return r         
        
    def findTagSeekPosition(self):
        self.ebml.tags = MatroskaTagsSeek
        self.ebml.quitOn = MatroskaTagsSeekQuit
        self.ebml.storePosition = True

        # get seek information
        all = self.ebml.parse()

        # store position and size information for later
        self.segmentPosition, self.segmentSize = all['Segment'][0][1:]
        self.seekHeadPosition, self.seekHeadSize = all['Segment'][0][0]['SeekHead'][0][1:]
        self.seekVoidPosition, self.seekVoidSize = all['Segment'][0][0]['EbmlVoid'][0][1:]
        
        nodes = all['Segment'][0][0]['SeekHead'][0][0]
        
        ploc = 0
        seekpos = -1
        
        for x in nodes['Seek']:
            if (x[0]['SeekID'][0][0] == '\x12T\xc3g'):
                ploc, seekpos = x[0]['SeekPosition'][0][:2]
                break
                
        self.ebml.tags = MatroskaTagsPosSeek
        self.ebml.quitOn = MatroskaTagsPosQuit
        self.ebml.storePosition = True
        
        ret = 0
        
        if ploc > 0:
            nodes = self.ebml.parse(ploc)
            try:
                ret = nodes['Tags'][0][1]
            except (KeyError):
                nodes = self.ebml.parse()
                try:
                    ret = nodes['Segment'][0][0]['Tags'][0][1]
                except (KeyError):
                    ret = 0
        
        return ret, ploc, seekpos
          
    
    def getTags(self):
        self.tagSeekLoc = self.findTagSeekPosition()[0]
        
        self.ebml.tags = MatroskaTagsTagging
        self.ebml.quitOn = {}
        self.ebml.storePosition = False
                
        if self.tagSeekLoc > 0:
            nodes = self.ebml.parse(from_ = self.tagSeekLoc)
            try:
                return nodes['Tags'][0]
            except:
                return False
        else:
            try:
                return nodes['Segment'][0]['Tags'][0]
            except:
                return False
            
class MatroskaTagReaderWriter(MatroskaTagReader):
    """
    Ebml Matroska Tag Writer
    
    usage: MatroskaTagReaderWriter(location).setTags(tags)
    usage: MatroskaTagReaderWriter(location).setTag(path, value)
    """

    def __init__(self, location):
        self.ebml = EbmlWriter(location)
        self.info = None

    def setTags(self, tags):
        self.tags = tags
        self.validated = False
        
    def _validateTag(self, name, tag):
        try:
            isMulti, parents = MatroskaTagValidity[name]
        except (KeyError), e:
            print >>sys.stderr, 'The following group is invalid: %s' % name, e
            return False
        
        # check if this tag is repeated too many times
        if (not isMulti) and (len(tag) > 1):
            print >>sys.stderr, 'The Tag: %s is repeated too many times (%d)' % name, len(tag)
            return False
        
        r = True
        
        # check all of the children
        for t in tag:
            try:
                for n, v in t.items():
                    # check if this tag is in the parents of the child
                    if name in MatroskaTagValidity[n][1]:
                        r = r and self._validateTag(n,v)
                    else:
                        print >>sys.stderr, 'The Tag: %s has invalid parent: %s' % n, name
                        return False
            except:
                return True
        
        return r
        
    def validateTags(self, tags = None):
        if tags != None:
            self.tags = tags
            
        # check if there are tags
        if self.tags == None:
            return False
        
        # validation stages:
        # 1: Loop through each Group, ensure they have the correct parent, are not mulitplied, etc        
        # 2: check target: TrackUID, ChapterUID, EditionUID, AttachmentUID for existance
        
        # 1
        if not self._validateTag('Tags', [self.tags]):
            print >>sys.stderr, 'Tag Validation Failed'
            return False
    
        # 2
       
        # find all TrackUID, ChapterUID, EditionUID, AttachmentUID's and compare against existing values       
        uidNames = ['TrackUID', 'ChapterUID', 'EditionUID', 'AttachmentUID']
        
        uids = self.getUIDs(uidNames)
        
        for i in self.tags['Tag']:
            for n in uidNames:
                # ensure any key errors are in the tags not the UIDs
                try:
                    u = uids[n]
                except (KeyError):
                    u = []
                    
                # check all the target uids
                try:
                    for t in i['Targets'][0][n]:
                        if not (t in uids[n]):
                            print >>sys.stderr, 'Unable to find match for %s: %d' % n, t
                            
                            # remove this entry from the tags
                            i['Targets'][0][n].remove(t)
                            
                            # check if this is the only one left and remove the dictionary item
                            if len(i['Targets'][0][n]) == 0:
                                del i['Targets'][0][n]
                                
                except (KeyError):
                    # do nothing on key error
                    continue
                   
        self.validated = True 
        return True
        
    def writeTags(self, tags = None):                
        import os
        
        if tags != None:
            self.setTags(tags)        
            
        # check if we have some tags to write
        if self.tags == None:
            return False
        
        if self.validated == False:
            if self.validateTags() == False:
                return False
        
        self.tagSeekLoc = self.findTagSeekPosition()[0]

        # check segment location
        if self.segmentPosition == -1:
            return False
    
        # keep a track of the sizes
        self.ebml.seek(0, os.SEEK_END)
        size = self.ebml.tell()
        
        if self.tagSeekLoc > 0:
            # move to the start of the existing tags, and remove from the file
            self.ebml.seek(self.tagSeekLoc)
            self.ebml.truncate()
        
        # write out the new tags
        try:
            atags = [self.tags]
            self.writeEbml('Tags', atags)
        
            # update the size of segment, as it will had of changed
            newSegmentSize = self.segmentSize + self.ebml.tell() - size
        
            self.ebml.updateSize(self.segmentPosition, newSegmentSize)
            
            if self.tagSeekLoc == 0:
                # write Tag Seek information
                self.addSeekInfo('\x12T\xc3g', size)

        except (Exception), e:
            print >>sys.stderr, 'ERROR', e
            return False
        
        return True
                
    def writeEbml(self, name, tag):
        try:
            # get the type of tag
            id, type = MatroskaTagsIndex[name]
            for i in tag:
                if type is MASTER:
                    self.ebml.startContainer(id)                    
                    for n, v in i.items():
                        self.writeEbml(n, v)
                    self.ebml.endContainer()
                elif type in (SINT, UINT, DATE):
                    self.ebml.writeInt(id, i)
                elif type is FLOAT:
                    self.ebml.writeFloat(id, i)
                elif type is STRING:
                    self.ebml.writeString(id, i)
                elif type is UTF8:
                    self.ebml.writeUTF8(id, i)
                elif type is BINARY:
                    self.ebml.writeBinary(id, i)
        except (KeyError), e:
            # cannot find name
            print >>sys.stderr, 'Unable to find Key: %s' % name, e

    def findTagSeekPosition(self):
        r, ploc, seekpos = MatroskaTagReader.findTagSeekPosition(self)
        
        # if the seek location is different to the actual location of the start of the Tags, fix        
        if r > 0 and r != ploc and seekpos != -1:
            self.ebml.updateIntValue(seekpos, r)
            
        return r, r, seekpos       
        
        
    # write seek information for the tags
    def addSeekInfo(self, sid, position):
        # 50 for maximum of id,size,id,size,sid,id,size,position
        if (self.seekVoidPosition > 0) and (self.seekVoidSize > 32):
            nodes = [{
                'SeekID': [sid],
                'SeekPosition': [position]
            }]
            
            # keep a record of the current file position
            cpos = self.ebml.tell()
            startPos = self.seekVoidPosition
            
            # find the end position of the EbmlVoid
            self.ebml.seek(self.seekVoidPosition)
            self.ebml.readID()
            self.ebml.readSize()
            
            endPos = self.ebml.tell() + self.seekVoidSize
                                  
            # write the new Seek information
            self.ebml.seek(self.seekVoidPosition)
            self.writeEbml('Seek', nodes)            
            voidPos = self.ebml.tell()
            
            # update the size of the seek head
            self.ebml.updateSize(self.seekHeadPosition, self.seekHeadSize + voidPos - self.seekVoidPosition)
            
            # write a void to cover the remaining space
            self.ebml._writeID(MatroskaSeekIndex['EbmlVoid'][0])
            
            # write dummy size information to start with
            self.ebml._writeSize(endPos - self.ebml.tell())
            
            # update with the correct size information, this will always 
            # be smaller than the previous value, so don't need to worry 
            # about size sizes
            self.ebml.updateSize(voidPos, endPos - self.ebml.tell())
            
            # return the position back to where we started
            self.ebml.seek(cpos)         

# vi: et sts=4 sw=4 ts=4
